import os
import numpy as np
from osgeo import gdal, osr
from tqdm import tqdm


def read_tif(file_path):
    dataset = gdal.Open(file_path)
    data = dataset.ReadAsArray()
    geotransform = dataset.GetGeoTransform()
    return data, geotransform


# 写入tif文件
def write_tif(file_path, data, geotransform, nodata):
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    dataset = driver.Create(file_path, cols, rows, 1, gdal.GDT_Float32)
    dataset.SetGeoTransform(geotransform)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    dataset.SetProjection(prj.ExportToWkt())
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    del dataset
    print("写入成功：{}".format(file_path))


# 根据像元行列号计算经纬度
def pixel2coord(geo_transform, x, y):
    x_origin, pixel_width, x_rotation, y_origin, y_rotation, pixel_height = geo_transform
    x_geo = x_origin + pixel_width * x
    y_geo = y_origin + pixel_height * y
    return x_geo, y_geo


# 根据经纬度计算像元行列号
def coord2pixel(geo_transform, x_geo, y_geo):
    x = int((x_geo - geo_transform[0]) / geo_transform[1])
    y = int((y_geo - geo_transform[3]) / geo_transform[5])
    return x, y


# 获取低分辨率像元对应所有高分辨率像元值
def get_high_res_pixels(high_res_geo_transform, high_res_data, low_res_geo_transform, x_low_res_pixel, y_low_res_pixel, low_res_to_high_res_ratio):
    # 计算低分辨率影像像元在高分辨率影像中的范围
    x_low_res_geo, y_low_res_geo = pixel2coord(low_res_geo_transform, x_low_res_pixel, y_low_res_pixel)
    x_low_res_geo = x_low_res_geo + low_res_geo_transform[1] / 2  # 左右偏移
    y_low_res_geo = y_low_res_geo + low_res_geo_transform[5] / 2  # 上下偏移

    # 计算行列号范围
    high_res_x, high_res_y = coord2pixel(high_res_geo_transform, x_low_res_geo, y_low_res_geo)

    high_res_x_min = int(high_res_x - low_res_to_high_res_ratio / 2)
    high_res_x_max = int(high_res_x + low_res_to_high_res_ratio / 2)
    high_res_y_min = int(high_res_y - low_res_to_high_res_ratio / 2)
    high_res_y_max = int(high_res_y + low_res_to_high_res_ratio / 2)

    # 排除边缘像元的影响
    if high_res_x_min < 0 and high_res_x_max > 0:
        high_res_x_min = 0
    if high_res_x_min < high_res_data.shape[1] < high_res_x_max:
        high_res_x_max = high_res_data.shape[1]
    if high_res_y_min < 0 and high_res_y_max > 0:
        high_res_y_min = 0
    if high_res_y_min < high_res_data.shape[0] < high_res_y_max:
        high_res_y_max = high_res_data.shape[0]

    # 获取范围内的像元值
    if 0 <= high_res_x_min <= high_res_data.shape[1] and 0 <= high_res_y_min <= high_res_data.shape[0]:
        # print(high_res_x_min, high_res_y_min, high_res_x_max, high_res_y_max)
        high_res_pixels = high_res_data[high_res_y_min:high_res_y_max, high_res_x_min:high_res_x_max]
        return high_res_pixels.astype(np.float32)
    else:
        return []


def downscale_formula(sm_LR, atdi, atdi_mean):

    sm_LR[sm_LR == -9999] = np.nan
    atdi[atdi == -9999] = np.nan
    atdi[np.isinf(atdi)] = np.nan
    atdi_mean[np.isinf(atdi_mean)] = np.nan
    print(np.nanmin(sm_LR))
    print(np.nanmax(sm_LR))
    print(np.nanmin(atdi))
    print(np.nanmax(atdi))
    print(np.nanmin(atdi_mean))
    print(np.nanmax(atdi_mean))
    # sm_HR = (atdi) / (atdi_mean) * sm_LR
    sm_HR = (1 - atdi) / (1 - atdi_mean) * sm_LR
    sm_HR[np.isnan(sm_HR)] = -9999
    return sm_HR


def pixel_match(atdi_path, sm_path):
    # 读取高分辨率和低分辨率的影像数据
    atdi_data, atdi_geotrans = read_tif(atdi_path)

    sm_data, sm_geotrans = read_tif(sm_path)
    print("===========")
    print(sm_data.shape)
    print(atdi_data.shape)
    # 获取像元比例
    low_res_to_high_res_ratio = int(sm_geotrans[1] / atdi_geotrans[1])
    # 预设两个列表存每个cci像元对应的ati和evi数据
    ati_in_sm = []
    # print(ati_in_sm)
    # 遍历cci每个像元，获取该像元中的所有ati和evi数据（一个二维数组）
    for y in range(0, sm_data.shape[0]):
        for x in range(0, sm_data.shape[1]):
            ati_in_sm.append(get_high_res_pixels(atdi_geotrans, atdi_data, sm_geotrans, x, y, low_res_to_high_res_ratio))
    print(np.array(ati_in_sm).shape)
    ati_in_sm = np.array(ati_in_sm).reshape((sm_data.shape[0], sm_data.shape[1]))
    print(ati_in_sm.shape)
    # print(ati_in_sm[1, 1])
    # 将数据维度调整为和atdi一致
    sm_h = np.zeros_like(atdi_data)
    atdi_mean_h = np.zeros_like(atdi_data)
    for y in range(0, atdi_data.shape[0]):
        for x in range(0, atdi_data.shape[1]):
            x_geo, y_geo = pixel2coord(atdi_geotrans, x, y)
            x_lr, y_lr = coord2pixel(sm_geotrans, x_geo, y_geo)
            sm_h[y, x] = sm_data[y_lr, x_lr]
            ati = ati_in_sm[y_lr, x_lr]
            atdi_mean_h[y, x] = np.nanmean(ati[ati != -9999])
    return sm_h, atdi_mean_h, atdi_data


def main(atdi_path, sm_path, output_path):
    atdi_ls = []
    sm_ls = []
    basename = ".tif"
    for root, dirs, files in os.walk(atdi_path):
        atdi_ls = [os.path.join(root, file) for file in files if file.endswith(basename)]

    for root, dirs, files in os.walk(sm_path):
        sm_ls = [os.path.join(root, file) for file in files if file.endswith(basename)]
    print(atdi_ls)
    print(sm_ls)
    for i, atdi_file in enumerate(atdi_ls):
        # todo 这里为了好调试，我按照一一对应写的，按自己的文件要求改回去即可
        # for sm_file in sm_ls:
        #     if atdi_file.split('/')[-1] == sm_file.split('/')[-1]:
        output_image = output_path + "\\" + atdi_ls[i].split('/')[-1].split('.')[0] + ".tif"
        atdi_data, atdi_geotrans = read_tif(atdi_file)
        sm_LR, sm_geotrans = read_tif(sm_ls[i])
        print(sm_LR.shape)
        sm, atdi_mean, atdi = pixel_match(atdi_file, sm_ls[i])
        print(sm.shape)
        print(atdi_mean.shape)
        print(atdi.shape)

        sm_HR = downscale_formula(sm, atdi, atdi_mean)
        write_tif(output_image, sm_HR, atdi_geotrans, -9999)


if __name__ == '__main__':
    # 指定输入文件路径，假设都转为tif格式了
    # ATDI_PATH = r"E:\data\A02\ATDI/lai/"
    SMI_PATH = r"E:\data\A03\ATI\rec_cubic_albedo/"

    SM_PATH = r"E:\data\A03\SM\cut/"
    OUTPUT_PATH = r"E:\data\A_04\downscaling\ati"
    main(SMI_PATH, SM_PATH, OUTPUT_PATH)